package com.hubspot.singularity.s3downloader.server;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.Inject;
import com.google.inject.Provider;
import com.google.inject.name.Named;
import com.hubspot.deploy.S3Artifact;
import com.hubspot.mesos.JavaUtils;
import com.hubspot.singularity.runner.base.sentry.SingularityRunnerExceptionNotifier;
import com.hubspot.singularity.s3.base.ArtifactDownloadRequest;
import com.hubspot.singularity.s3.base.ArtifactManager;
import com.hubspot.singularity.s3downloader.SingularityS3DownloaderMetrics;
import com.hubspot.singularity.s3downloader.config.SingularityS3DownloaderConfiguration;
import com.hubspot.singularity.s3downloader.config.SingularityS3DownloaderModule;
import java.io.IOException;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import javax.servlet.http.HttpServletResponse;
import org.eclipse.jetty.continuation.Continuation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SingularityS3DownloaderCoordinator implements DownloadListener {
  private static final Logger LOG = LoggerFactory.getLogger(
    SingularityS3DownloaderCoordinator.class
  );

  private final SingularityS3DownloaderConfiguration configuration;
  private final SingularityS3DownloaderMetrics metrics;
  private final Provider<ArtifactManager> artifactManagerProvider;
  private final ConcurrentMap<S3Artifact, SingularityS3DownloaderAsyncHandler> downloadRequestToHandler;
  private final ScheduledThreadPoolExecutor downloadJoinerService;
  private final ThreadPoolExecutor downloadService;
  private final ListeningExecutorService listeningDownloadWrapper;
  private final ExecutorService listeningResponseExecutorService;
  private final SingularityRunnerExceptionNotifier exceptionNotifier;

  @Inject
  public SingularityS3DownloaderCoordinator(
    SingularityS3DownloaderConfiguration configuration,
    SingularityS3DownloaderMetrics metrics,
    Provider<ArtifactManager> artifactManagerProvider,
    @Named(
      SingularityS3DownloaderModule.ENQUEUE_EXECUTOR_SERVICE
    ) ScheduledThreadPoolExecutor downloadJoinerService,
    @Named(
      SingularityS3DownloaderModule.DOWNLOAD_EXECUTOR_SERVICE
    ) ThreadPoolExecutor downloadService,
    SingularityRunnerExceptionNotifier exceptionNotifier
  ) {
    this.configuration = configuration;
    this.metrics = metrics;
    this.artifactManagerProvider = artifactManagerProvider;
    this.downloadJoinerService = downloadJoinerService;
    this.downloadService = downloadService;
    this.downloadRequestToHandler = Maps.newConcurrentMap();
    this.listeningDownloadWrapper = MoreExecutors.listeningDecorator(downloadService);
    this.listeningResponseExecutorService = Executors.newCachedThreadPool();
    this.exceptionNotifier = exceptionNotifier;
  }

  public void register(
    final Continuation continuation,
    final ArtifactDownloadRequest artifactDownloadRequest
  ) {
    final DownloadJoiner downloadJoiner = new DownloadJoiner(
      continuation,
      artifactDownloadRequest
    );

    downloadJoinerService.submit(downloadJoiner);
  }

  @Override
  public void notifyDownloadFinished(SingularityS3DownloaderAsyncHandler handler) {
    boolean removed = downloadRequestToHandler.remove(handler.getS3Artifact(), handler);
    LOG.debug(
      "Handler for artifact {} finished download - removed {}",
      handler.getS3Artifact(),
      removed
    );
  }

  private class DownloadJoiner implements Runnable {
    private final long start;
    private final Continuation continuation;
    private final ArtifactDownloadRequest artifactDownloadRequest;

    private DownloadJoiner(
      Continuation continuation,
      ArtifactDownloadRequest artifactDownloadRequest
    ) {
      this.continuation = continuation;
      this.artifactDownloadRequest = artifactDownloadRequest;
      this.start = System.currentTimeMillis();
    }

    private void reEnqueue() {
      LOG.debug(
        "Re-enqueueing request for {}, waiting {}, ({} active, {} queue, {} max), total time {}",
        artifactDownloadRequest.getTargetDirectory(),
        JavaUtils.durationFromMillis(configuration.getMillisToWaitForReEnqueue()),
        downloadJoinerService.getActiveCount(),
        downloadJoinerService.getQueue().size(),
        configuration.getNumEnqueueThreads(),
        JavaUtils.duration(start)
      );

      downloadJoinerService.schedule(
        this,
        configuration.getMillisToWaitForReEnqueue(),
        TimeUnit.MILLISECONDS
      );
    }

    private boolean addDownloadRequest() {
      SingularityS3DownloaderAsyncHandler existingHandler = downloadRequestToHandler.get(
        artifactDownloadRequest.getS3Artifact()
      );

      if (existingHandler != null) {
        return false;
      }

      final SingularityS3DownloaderAsyncHandler newHandler = new SingularityS3DownloaderAsyncHandler(
        artifactManagerProvider.get(),
        artifactDownloadRequest,
        continuation,
        metrics,
        exceptionNotifier,
        SingularityS3DownloaderCoordinator.this
      );

      existingHandler =
        downloadRequestToHandler.putIfAbsent(
          artifactDownloadRequest.getS3Artifact(),
          newHandler
        );

      if (existingHandler != null) {
        return false;
      }

      LOG.trace(
        "Queing new downloader for {} ({} handlers, {} active threads, {} queue size, {} max) after {}",
        artifactDownloadRequest,
        downloadRequestToHandler.size(),
        downloadService.getActiveCount(),
        downloadService.getQueue().size(),
        configuration.getNumDownloaderThreads(),
        JavaUtils.duration(start)
      );

      ListenableFuture<?> future = listeningDownloadWrapper.submit(newHandler);

      future.addListener(
        new Runnable() {

          @Override
          public void run() {
            notifyDownloadFinished(newHandler);
          }
        },
        listeningResponseExecutorService
      );

      return true;
    }

    private void handleContinuationExpired() {
      try {
        LOG.info(
          "Continuation expired for {} after {} - returning 500",
          artifactDownloadRequest,
          JavaUtils.duration(start)
        );
        ((HttpServletResponse) continuation.getServletResponse()).sendError(
            500,
            "Hit client timeout"
          );
      } catch (Throwable t) {
        LOG.warn(
          "{} while sending error after continuation for {}",
          t.getClass().getSimpleName(),
          artifactDownloadRequest.getTargetDirectory()
        );
      } finally {
        continuation.complete();
      }
    }

    @Override
    public void run() {
      try {
        if (!addDownloadRequest()) {
          if (continuation.isExpired()) {
            handleContinuationExpired();
            return;
          }
          reEnqueue();
        }
      } catch (Throwable t) {
        LOG.error(
          "While trying to enqueue {}",
          artifactDownloadRequest.getTargetDirectory(),
          t
        );
        exceptionNotifier.notify(
          String.format("Error enqueuing download (%s)", t.getMessage()),
          t,
          ImmutableMap.of("targetDirectory", artifactDownloadRequest.getTargetDirectory())
        );
        try {
          ((HttpServletResponse) continuation.getServletResponse()).sendError(500);
        } catch (IOException e) {
          LOG.error(
            "Couldn't send error for {}",
            artifactDownloadRequest.getTargetDirectory(),
            e
          );
        } finally {
          continuation.complete();
        }
      }
    }
  }
}
